#!/usr/bin/python

__author__ = "Jimmy Saw"
__copyright__ = "Copyright 2011, extracts intergenic sequences and BLAST them to get hits"
__credits__ = "Jimmy Saw"
__email__ = "jimmysaw@gmail.com"

"""
TO DO: 
Usage: python identify_igs.py annofile.txt seqfile.fasta prefix

"""

import sys
from Bio import SeqIO
from Bio.Blast import NCBIWWW
from Bio.Blast import NCBIXML

annofile = sys.argv[1]
seqfile = sys.argv[2]
prefix = sys.argv[3]
af = open(annofile, "rU")
sf = open(seqfile, "rU")
rec = SeqIO.read(sf, "fasta")
lines = af.readlines()
num = len(lines)

i = 0

while i < num:
    if i == 0:
        curr_line = lines[i].split('\t')
        curr_id = curr_line[0]
        curr_start = int(curr_line[3])
        substop = curr_start - 1
        igs_id = prefix + "_IGS_" + str(i).zfill(4) + "_" + "0-" + str(substop)
        subseq = rec.seq[0:substop]
        if len(subseq) > 90:
            print "Running BLASTx of " + igs_id
            result_handle = NCBIWWW.qblast("blastx", "nr", subseq, expect=0.00001, filter=None)
            blast_rec = NCBIXML.read(result_handle)
            if blast_rec.alignments != None:
                for alignment in blast_rec.alignments:
                    for hsp in alignment.hsps:
                        print alignment.length
            save_file = open(igs_id + ".xml", "w")
            save_file.write(result_handle.read()) 
            save_file.close()
            result_handle.close()
            print "Done BLASTx"

    if i == num - 1:
        curr_line = lines[i].split('\t')
        curr_stop = int(curr_line[4])
        substart = curr_stop + 1
        substop = len(rec.seq)
        igs_id = prefix + "_IGS_" + str(i).zfill(4) + "_" + str(substart) + "-" + str(substop)
        subseq = rec.seq[substart:substop]
        if len(subseq) > 90:
            print "Running BLASTx of " + igs_id
            result_handle = NCBIWWW.qblast("blastx", "nr", subseq, expect=0.00001, filter=None)
            if blast_rec.alignments != None:
                for alignment in blast_rec.alignments:
                    for hsp in alignment.hsps:
                        print alignment.length
            save_file = open(igs_id + ".xml", "w")
            save_file.write(result_handle.read()) 
            save_file.close()
            result_handle.close()
            print "Done BLASTx"

    if i != num - 1 and i != 0:
        curr_line = lines[i].split('\t')
        curr_id = curr_line[0]
        curr_locus_tag = curr_line[1]
        curr_feat_type = curr_line[2]
        curr_start = int(curr_line[3])
        curr_stop = int(curr_line[4])
        curr_frame = curr_line[5]

        next_line = lines[i + 1].split('\t')
        next_id = next_line[0]
        next_locus_tag = next_line[1]
        next_feat_type = next_line[2]
        next_start = int(next_line[3])
        next_stop = int(next_line[4])
        next_frame = next_line[5]

        prev_line = lines[i - 1].split('\t')
        prev_id = prev_line[0]
        prev_locus_tag = prev_line[1]
        prev_feat_type = prev_line[2]
        prev_start = int(prev_line[3])
        prev_stop = int(prev_line[4])
        prev_frame = prev_line[5]

        if curr_feat_type == "CDS":
            substart = prev_stop + 1
            substop = curr_start - 1
            if substop > substart:
                igs_id = prefix + "_IGS_" + str(i).zfill(4) + "_" + str(substart) + "-" + str(substop)
                subseq = rec.seq[substart:substop]
                if len(subseq) > 90:
                    print "Running BLASTx of " + igs_id
                    result_handle = NCBIWWW.qblast("blastx", "nr", subseq, expect=0.00001, filter=None)
                    if blast_rec.alignments != None:
                        for alignment in blast_rec.alignments:
                            for hsp in alignment.hsps:
                                print alignment.length
                    save_file = open(igs_id + ".xml", "w")
                    save_file.write(result_handle.read()) 
                    save_file.close()
                    result_handle.close()
                    print "Done BLASTx"
    i += 1



af.close()
sf.close()
